{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "from miasm.analysis.machine import Machine\n",
    "from miasm.ir.symbexec import SymbolicExecutionEngine\n",
    "from miasm.ir.ir import IRCFG\n",
    "from miasm.expression.expression import LocKey\n",
    "from miasm.arch.x86.regs import *\n",
    "from miasm.core import asmblock\n",
    "from miasm.core.locationdb import LocationDB\n",
    "from miasm.analysis.binary import Container\n",
    "from future.utils import viewitems\n",
    "from miasm.ir.translators.translator import Translator\n",
    "import networkx as nx\n",
    "import random\n",
    "import z3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip -n -P infected vipasana.zip\n",
    "!unzip -n -P infected asprox.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = 'vipasana.bin'\n",
    "target_addr = 0x434DF0\n",
    "#filename = '../asprox.bin'\n",
    "#target_addr = 0x100091AC\n",
    "idc = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def syntax_compare(blk0, blk1):\n",
    "    if len(blk0.lines) != len(blk1.lines):\n",
    "        return False\n",
    "\n",
    "    for l0, l1 in zip(blk0.lines, blk1.lines):\n",
    "        if str(l0)[0] == 'J':\n",
    "            instr0 = str(l0).split(' ')[0]\n",
    "            instr1 = str(l1).split(' ')[0]\n",
    "            if instr0 != instr1:\n",
    "                return False\n",
    "        else:\n",
    "            if str(l0) != str(l1):\n",
    "                return False\n",
    "\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def semantic_compare(blk0, blk1, ir_arch0, ir_arch1, asmcfg, flag_cmp=False):\n",
    "    src_ircfg = IRCFG(None, ir_arch0.loc_db)\n",
    "    try:\n",
    "        ir_arch0.add_asmblock_to_ircfg(blk0, src_ircfg)\n",
    "    except NotImplementedError:\n",
    "        return False\n",
    "\n",
    "    dst_ircfg = IRCFG(None, ir_arch1.loc_db)\n",
    "    try:\n",
    "        ir_arch1.add_asmblock_to_ircfg(blk1, dst_ircfg)\n",
    "    except NotImplementedError:\n",
    "        return False\n",
    "\n",
    "    if len(src_ircfg.blocks) != len(dst_ircfg.blocks):\n",
    "        return False\n",
    "\n",
    "    for src_lbl, dst_lbl in zip(src_ircfg.blocks, dst_ircfg.blocks):\n",
    "\n",
    "        src_irb = src_ircfg.blocks.get(src_lbl, None)\n",
    "        dst_irb = dst_ircfg.blocks.get(dst_lbl, None)\n",
    "\n",
    "        r = execute_symbolic_execution(\n",
    "                            src_irb, dst_irb, \n",
    "                            ir_arch0, ir_arch1, \n",
    "                            src_ircfg, dst_ircfg,\n",
    "                            flag_cmp)\n",
    "        if r is False:\n",
    "            return False\n",
    "\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute_symbolic_execution(src_irb, dst_irb, \n",
    "                                ir_arch0, ir_arch1, \n",
    "                                src_ircfg, dst_ircfg,\n",
    "                                flag_cmp):\n",
    "\n",
    "    # Ready for Symbolic Execution\n",
    "    src_symbols = {}\n",
    "    dst_symbols = {}\n",
    "\n",
    "    # regs\n",
    "    for i, r in enumerate(all_regs_ids):\n",
    "        src_symbols[r] = all_regs_ids_init[i]\n",
    "        dst_symbols[r] = all_regs_ids_init[i]\n",
    "\n",
    "\n",
    "    # Run symbolic execution\n",
    "    src_sb = SymbolicExecutionEngine(ir_arch0, src_symbols)\n",
    "\n",
    "    for assignblk in src_irb:\n",
    "        skip_update = False\n",
    "        for dst, src in viewitems(assignblk):\n",
    "            if str(dst) in ['EIP', 'IRDst']:\n",
    "                skip_update = True\n",
    "\n",
    "        if not skip_update:\n",
    "            src_sb.eval_updt_assignblk(assignblk)\n",
    "\n",
    "    dst_sb = SymbolicExecutionEngine(ir_arch1, dst_symbols)\n",
    "\n",
    "    for assignblk in dst_irb:\n",
    "        skip_update = False\n",
    "        for dst, src in viewitems(assignblk):\n",
    "            if str(dst) in ['EIP', 'IRDst']:\n",
    "                skip_update = True\n",
    "\n",
    "        if not skip_update:\n",
    "            dst_sb.eval_updt_assignblk(assignblk)\n",
    "\n",
    "    # Equivalence Checking\n",
    "\n",
    "    src_sb.del_mem_above_stack(ir_arch0.sp)\n",
    "    dst_sb.del_mem_above_stack(ir_arch1.sp)\n",
    "\n",
    "    all_memory_ids  = [k for k, v in dst_sb.symbols.memory()] + [k for k, v in src_sb.symbols.memory()]\n",
    "\n",
    "    for k in all_regs_ids + all_memory_ids:\n",
    "\n",
    "        if str(k) == 'EIP':\n",
    "            continue\n",
    "\n",
    "        if not flag_cmp and k in [zf, nf, pf, of, cf, af, df, tf]:\n",
    "            continue\n",
    "\n",
    "        v0 = src_sb.symbols[k]\n",
    "        v1 = dst_sb.symbols[k]\n",
    "\n",
    "        if v0 == v1:\n",
    "            continue\n",
    "\n",
    "        solver = z3.Solver()\n",
    "        try:\n",
    "            z3_r_cond = Translator.to_language('z3').from_expr(v0)\n",
    "        except NotImplementedError:\n",
    "            return False\n",
    "\n",
    "        try:\n",
    "            z3_l_cond = Translator.to_language('z3').from_expr(v1)\n",
    "        except NotImplementedError:\n",
    "            return False\n",
    "\n",
    "        # HERE\n",
    "        solver.add(z3.Not(z3_r_cond == z3_l_cond))\n",
    "\n",
    "        r = solver.check()\n",
    "        if r == z3.unsat:\n",
    "            continue\n",
    "\n",
    "        else:\n",
    "            #print(solver.model()) # Counterexample\n",
    "            return False\n",
    "\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc_db = LocationDB()\n",
    "with open(filename, 'rb') as fstream:                                      \n",
    "    cont = Container.from_stream(fstream, loc_db)\n",
    "    \n",
    "machine = Machine('x86_32')\n",
    "mdis = machine.dis_engine(cont.bin_stream, loc_db=cont.loc_db)\n",
    "ir_arch0 = machine.ira(mdis.loc_db)\n",
    "ir_arch1 = machine.ira(mdis.loc_db)\n",
    "\n",
    "asmcfg = mdis.dis_multiblock(target_addr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_blocks = []\n",
    "for cn, block in enumerate(asmcfg.blocks):\n",
    "    target_blocks.append(block)\n",
    "\n",
    "results = {}\n",
    "\n",
    "for src_blk in target_blocks:\n",
    "    src_ldl = src_blk._loc_key\n",
    "\n",
    "    # Skip a basic block containing only single instruction\n",
    "    if len(src_blk.lines) == 1 and src_blk.lines[0].dstflow():\n",
    "        continue\n",
    "\n",
    "    for dst_blk in target_blocks:\n",
    "        dst_ldl = dst_blk._loc_key\n",
    "\n",
    "        # Skip a basic block containing only single instruction\n",
    "        if len(dst_blk.lines) == 1 and dst_blk.lines[0].dstflow():\n",
    "            continue\n",
    "\n",
    "        if src_ldl == dst_ldl:\n",
    "            continue\n",
    "\n",
    "        if (src_ldl, dst_ldl) in results.keys() or \\\n",
    "            (dst_ldl, src_ldl) in results.keys():\n",
    "            continue\n",
    "        \n",
    "        r_syntax = syntax_compare(src_blk, dst_blk)\n",
    "\n",
    "        if r_syntax:\n",
    "            # If the syntax of two blocks is same, then the semantics of them is also same.\n",
    "            r_semantic = True\n",
    "        else:\n",
    "            # Otherwise, need to compare the semantics of them\n",
    "            r_semantic = semantic_compare(src_blk, dst_blk, ir_arch0, ir_arch1, asmcfg)\n",
    "\n",
    "        results[(src_ldl, dst_ldl)] = [(r_syntax, r_semantic)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if idc:\n",
    "    G = nx.Graph()\n",
    "    G.add_nodes_from(target_blocks)\n",
    "\n",
    "    for k, v in viewitems(results):\n",
    "        if v[0][0] or v[0][1]:\n",
    "            G.add_edge(k[0], k[1])\n",
    "\n",
    "    # Return a list containing randomlly generated colors\n",
    "    def gen_random_color():\n",
    "\n",
    "        ret = []\n",
    "\n",
    "        r = [x for x in range(256)]\n",
    "        g = [x for x in range(256)]\n",
    "        b = [x for x in range(256)]\n",
    "        random.shuffle(r)\n",
    "        random.shuffle(g)\n",
    "        random.shuffle(b)\n",
    "\n",
    "        for a2, a1, a0 in zip(r,g,b):\n",
    "            color = a2 << 16 | a1 << 8 | a0\n",
    "            ret.append(color)\n",
    "\n",
    "        return ret\n",
    "\n",
    "    random_colors = gen_random_color()\n",
    "    body = ''\n",
    "\n",
    "    for n, conn_nodes in enumerate(nx.connected_components(G)):\n",
    "\n",
    "        if len(conn_nodes) == 1:\n",
    "            continue\n",
    "\n",
    "        for node in conn_nodes: # node is asmblk\n",
    "\n",
    "            if isinstance(node, LocKey):\n",
    "                asmblk = asmcfg.loc_key_to_block(node)\n",
    "                if asmblk:\n",
    "                    for l in asmblk.lines:\n",
    "                        body += 'SetColor(0x%08x, CIC_ITEM, 0x%x);\\n'%(l.offset, random_colors[n])\n",
    "            else:\n",
    "                for l in node.lines:\n",
    "                    body += 'SetColor(0x%08x, CIC_ITEM, 0x%x);\\n'%(l.offset, random_colors[n])\n",
    "        \n",
    "    header = '''\n",
    "#include <idc.idc>\n",
    "static main()\n",
    "{\n",
    "'''\n",
    "    footer = '''\n",
    "}\n",
    "'''\n",
    "\n",
    "    f = open('eq-color.idc', 'w')\n",
    "    f.write(header+body+footer)\n",
    "    f.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
